import time
import datetime
import json
import threading
from services.database_service import DatabaseService
from services.chroma_service import ChromaService
from models.sync_config import SyncTask, SyncTaskRun, db
from llm.model_factory import ModelFactory
from models.document import Knowledge
from utils.tools import deterministic_uuid
from flask import current_app
from llm.model_registry import ModelRegistry
from llm.model_service import ModelService

model_client = ModelFactory.create_model()
chroma_service = ChromaService()

class SyncService:
    """同步服务，用于从关系型数据库同步数据到ChromaDB"""
    
    # 存储正在运行的同步任务
    _running_tasks = {}
    _lock = threading.Lock()
    
    @staticmethod
    def start_sync_task(sync_task):
        """启动同步任务"""
        with SyncService._lock:
            # 检查任务是否已在运行
            if sync_task.id in SyncService._running_tasks:
                return False, "任务已在运行中"
            
            # 创建任务运行记录
            task_run = SyncTaskRun(
                sync_task_id=sync_task.id,
                status='running'
            )
            db.session.add(task_run)
            db.session.commit()
            
            # 获取当前应用实例
            app = current_app._get_current_object()
            
            # 启动任务线程
            thread = threading.Thread(
                target=SyncService._sync_task_thread,
                args=(sync_task.id, task_run.id, app)
            )
            thread.daemon = True
            thread.start()
            
            # 将任务添加到运行列表
            SyncService._running_tasks[sync_task.id] = {
                'thread': thread,
                'task_run_id': task_run.id
            }
            
            return True, task_run.id
    
    @staticmethod
    def _sync_task_thread(task_id, run_id, app):
        """同步任务线程"""
        with app.app_context():
            sync_task = SyncTask.query.get(task_id)
            task_run = SyncTaskRun.query.get(run_id)
            
            try:
                # 获取数据库连接
                db_connection = sync_task.database_connection
                
                # 获取总记录数
                success, total_count = DatabaseService.count_query_results(
                    db_connection, 
                    sync_task.sql_query
                )
                
                if not success:
                    raise Exception(f"计算总记录数失败: {total_count}")
                
                # 更新任务运行状态
                task_run.total_records = total_count
                db.session.commit()
                
                # 解析字段配置
                document_fields = json.loads(sync_task.document_fields)
                metadata_fields = json.loads(sync_task.metadata_fields)
                embedding_fields = json.loads(sync_task.embedding_fields)

                knowledge = Knowledge.query.get(sync_task.knowledge_id)
                if not knowledge:
                    raise Exception(f"知识库不存在")
                
                collection_name = knowledge.collection_name
                embedding_model = knowledge.embedding_model
                
                # 分批处理数据
                processed_count = 0
                start_time = time.time()
                
                # 流式查询处理
                for batch in DatabaseService.stream_query_results(db_connection, sync_task.sql_query):
                    if not batch:
                        continue
                    
                    # 处理每条记录
                    documents = []
                    metadatas = []
                    ids = []
                    texts_for_embedding = []
                    
                    for record in batch:
                        # 准备document
                        if document_fields:
                            doc_dict = {field: str(record.get(field, "")) for field in document_fields}
                            doc_content = json.dumps(doc_dict, ensure_ascii=False)
                        else:
                            # 如果没有指定字段，使用所有字段
                            doc_dict = {k: str(v) for k, v in record.items()}
                            doc_content = json.dumps(doc_dict, ensure_ascii=False)
                        
                        # 准备metadata
                        metadata = {}
                        if metadata_fields:
                            for field in metadata_fields:
                                metadata[field] = str(record.get(field, ""))
                        
                        # 生成ID
                        record_id = deterministic_uuid(doc_content)
                        
                        documents.append(doc_content)
                        metadatas.append(metadata)
                        ids.append(record_id)
                        
                        # 准备embedding文本
                        if embedding_fields:
                            embedding_text = " ".join([str(record.get(field, "")) for field in embedding_fields])
                        else:
                            embedding_text = doc_content
                        
                        texts_for_embedding.append(embedding_text)
                    
                    # 使用模型生成向量嵌入
                    model_tuple = ModelRegistry.split_model_name(embedding_model)
                    embed_provider = model_tuple[0]
                    embed_model = model_tuple[1]
                    embeddings = ModelService.call(
                        provider_name=embed_provider, 
                        model_name=embed_model, 
                        mode="embed", stream=False, text=texts_for_embedding)
                    # embeddings = model_client.get_embeddings(texts_for_embedding, embedding_model)
                    
                    # 添加到ChromaDB
                    chroma_service.add_records(
                        collection_name=collection_name,
                        embeddings=embeddings,
                        documents=documents,
                        metadatas=metadatas,
                        ids=ids
                    )
                    
                    # 更新进度
                    processed_count += len(batch)
                    task_run.processed_records = processed_count
                    
                    # 计算剩余时间
                    elapsed_time = time.time() - start_time
                    if processed_count > 0:
                        avg_time_per_record = elapsed_time / processed_count
                        remaining_records = total_count - processed_count
                        remaining_time = avg_time_per_record * remaining_records
                    
                    db.session.commit()
                
                # 更新任务状态
                task_run.status = 'completed'
                task_run.end_time = datetime.datetime.now()
                sync_task.last_run_at = datetime.datetime.now()
                sync_task.status = 'completed'
                db.session.commit()
            
            except Exception as e:
                # 处理异常
                task_run.status = 'failed'
                task_run.end_time = datetime.datetime.now()
                task_run.error_message = str(e)
                sync_task.status = 'failed'
                db.session.commit()
            
            finally:
                # 从运行列表中移除任务
                with SyncService._lock:
                    if task_id in SyncService._running_tasks:
                        del SyncService._running_tasks[task_id]
    
    @staticmethod
    def get_task_status(task_run_id):
        """获取任务运行状态"""
        task_run = SyncTaskRun.query.get(task_run_id)
        if not task_run:
            return None
        
        return task_run.to_dict()
    
    @staticmethod
    def cancel_task(task_id):
        """取消任务"""
        with SyncService._lock:
            if task_id not in SyncService._running_tasks:
                return False, "任务未在运行"
            
            # 更新任务状态
            task_run_id = SyncService._running_tasks[task_id]['task_run_id']
            task_run = SyncTaskRun.query.get(task_run_id)
            if task_run:
                task_run.status = 'cancelled'
                task_run.end_time = datetime.datetime.now()
                db.session.commit()
            
            # 移除任务
            del SyncService._running_tasks[task_id]
            
            return True, "任务已取消"